package tenant

import (
	"fmt"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
	"strings"
)

// TenantConditionEnhancer 增强版租户条件添加器
type TenantConditionEnhancer struct {
	targetTables map[string]bool
	modified     bool
}

// NewTenantConditionEnhancer 创建新的增强版租户条件添加器
func NewTenantConditionEnhancer(targetTables []string) *TenantConditionEnhancer {
	tableMap := make(map[string]bool)
	for _, table := range targetTables {
		tableMap[table] = true
	}
	return &TenantConditionEnhancer{
		targetTables: tableMap,
		modified:     false,
	}
}

// Enter 进入节点时的处理
func (v *TenantConditionEnhancer) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
	switch node := n.(type) {
	case *ast.SelectStmt:
		v.processSelectStmt(node)
		return node, false
	case *ast.UpdateStmt:
		v.processUpdateStmt(node)
		return node, false
	case *ast.DeleteStmt:
		v.processDeleteStmt(node)
		return node, false
	}
	return n, false
}

// Leave 离开节点时的处理
func (v *TenantConditionEnhancer) Leave(n ast.Node) (node ast.Node, ok bool) {
	return n, true
}

// processSelectStmt 处理SELECT语句
func (v *TenantConditionEnhancer) processSelectStmt(stmt *ast.SelectStmt) {
	if stmt.From != nil {
		hasTargetTable := v.checkTableRefs(stmt.From.TableRefs)
		if hasTargetTable {
			v.addTenantConditionToSelect(stmt)
		}
	}
}

// processUpdateStmt 处理UPDATE语句
func (v *TenantConditionEnhancer) processUpdateStmt(stmt *ast.UpdateStmt) {
	if stmt.TableRefs != nil {
		hasTargetTable := v.checkTableRefs(stmt.TableRefs.TableRefs)
		if hasTargetTable {
			v.addTenantConditionToUpdate(stmt)
		}
	}
}

// processDeleteStmt 处理DELETE语句
func (v *TenantConditionEnhancer) processDeleteStmt(stmt *ast.DeleteStmt) {
	if stmt.TableRefs != nil {
		hasTargetTable := v.checkTableRefs(stmt.TableRefs.TableRefs)
		if hasTargetTable {
			v.addTenantConditionToDelete(stmt)
		}
	}
}

// checkTableRefs 检查表引用中是否包含目标表
func (v *TenantConditionEnhancer) checkTableRefs(tableRefs *ast.Join) bool {
	if tableRefs == nil {
		return false
	}

	// 检查左表
	if leftTable, ok := tableRefs.Left.(*ast.TableSource); ok {
		if tableName, ok := leftTable.Source.(*ast.TableName); ok {
			if v.targetTables[tableName.Name.String()] {
				return true
			}
		}
	}

	// 检查右表
	if tableRefs.Right != nil {
		if rightTable, ok := tableRefs.Right.(*ast.TableSource); ok {
			if tableName, ok := rightTable.Source.(*ast.TableName); ok {
				if v.targetTables[tableName.Name.String()] {
					return true
				}
			}
		}
	}

	return false
}

// addTenantConditionToSelect 为SELECT语句添加tenant_id条件
func (v *TenantConditionEnhancer) addTenantConditionToSelect(stmt *ast.SelectStmt) {
	tenantCondition := v.createTenantCondition()

	if stmt.Where != nil {
		// 如果已有WHERE条件，用AND连接
		stmt.Where = &ast.BinaryOperationExpr{
			Op: ast.BinaryOpAnd,
			L:  stmt.Where,
			R:  tenantCondition,
		}
	} else {
		// 如果没有WHERE条件，直接设置
		stmt.Where = tenantCondition
	}
	v.modified = true
	fmt.Printf("为SELECT语句添加tenant_id条件\n")
}

// addTenantConditionToUpdate 为UPDATE语句添加tenant_id条件
func (v *TenantConditionEnhancer) addTenantConditionToUpdate(stmt *ast.UpdateStmt) {
	tenantCondition := v.createTenantCondition()

	if stmt.Where != nil {
		// 如果已有WHERE条件，用AND连接
		stmt.Where = &ast.BinaryOperationExpr{
			Op: ast.BinaryOpAnd,
			L:  stmt.Where,
			R:  tenantCondition,
		}
	} else {
		// 如果没有WHERE条件，直接设置
		stmt.Where = tenantCondition
	}
	v.modified = true
	fmt.Printf("为UPDATE语句添加tenant_id条件\n")
}

// addTenantConditionToDelete 为DELETE语句添加tenant_id条件
func (v *TenantConditionEnhancer) addTenantConditionToDelete(stmt *ast.DeleteStmt) {
	tenantCondition := v.createTenantCondition()

	if stmt.Where != nil {
		// 如果已有WHERE条件，用AND连接
		stmt.Where = &ast.BinaryOperationExpr{
			Op: ast.BinaryOpAnd,
			L:  stmt.Where,
			R:  tenantCondition,
		}
	} else {
		// 如果没有WHERE条件，直接设置
		stmt.Where = tenantCondition
	}
	v.modified = true
	fmt.Printf("为DELETE语句添加tenant_id条件\n")
}

// createTenantCondition 创建tenant_id条件
func (v *TenantConditionEnhancer) createTenantCondition() ast.ExprNode {
	return &ast.BinaryOperationExpr{
		Op: ast.BinaryOpEQ,
		L:  &ast.ColumnNameExpr{Name: &ast.ColumnName{Name: ast.NewCIStr("tenant_id")}},
		R:  &ast.ParamMarkerExpr{},
	}
}

// AddTenantConditionToSQLEnhanced 增强版SQL添加tenant_id条件
func AddTenantConditionToSQLEnhanced(sql string, targetTables []string) (string, error) {
	p := parser.New()

	// 解析SQL语句
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return "", fmt.Errorf("解析SQL失败: %v", err)
	}

	// 创建访问者
	enhancer := NewTenantConditionEnhancer(targetTables)

	// 遍历所有语句节点
	for _, stmt := range stmtNodes {
		stmt.Accept(enhancer)
	}

	// 将修改后的AST转换回SQL字符串
	restorer := NewSQLRestorer()
	var resultSQL string

	if len(stmtNodes) == 1 {
		// 单个语句
		resultSQL, err = restorer.RestoreSQL(stmtNodes[0])
		if err != nil {
			return "", err
		}
	} else {
		// 多个语句，用分号连接
		var sqls []string
		for _, stmt := range stmtNodes {
			sqlStr, err := restorer.RestoreSQL(stmt)
			if err != nil {
				return "", err
			}
			sqls = append(sqls, sqlStr)
		}
		resultSQL = strings.Join(sqls, "; ")
	}

	return resultSQL, nil
}

// ParseAndAddTenantConditionEnhanced 增强版解析SQL并添加tenant_id条件的示例函数
func ParseAndAddTenantConditionEnhanced() {
	// 示例SQL语句
	sqlExamples := []string{
		"SELECT * FROM t_user WHERE id = 1",
		"SELECT * FROM t_dept WHERE name = 'IT'",
		"SELECT u.name, d.name FROM t_user u JOIN t_dept d ON u.dept_id = d.id",
		"UPDATE t_user SET name = 'John' WHERE id = 1",
		"DELETE FROM t_dept WHERE id = 1",
		"SELECT * FROM t_user WHERE age > 18 AND status = 'active'",
		"SELECT * FROM other_table WHERE id = 1", // 这个表不在目标列表中
		"SELECT * FROM t_user",                   // 没有WHERE条件的情况
	}

	// 目标表名列表
	targetTables := []string{"t_user", "t_dept"}

	fmt.Println("=== SQL添加tenant_id条件示例 (增强版) ===")

	for i, sql := range sqlExamples {
		fmt.Printf("\n--- 示例 %d ---\n", i+1)
		fmt.Printf("原始SQL: %s\n", sql)

		// 解析并添加tenant_id条件
		modifiedSQL, err := AddTenantConditionToSQLEnhanced(sql, targetTables)
		if err != nil {
			fmt.Printf("处理失败: %v\n", err)
			continue
		}

		fmt.Printf("修改后SQL: %s\n", modifiedSQL)
		fmt.Println("处理完成")
	}
}
